# MIT License
# 
# Copyright (c) 2024 Jianbo ZHU
# 
# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:
# 
# The above copyright notice and this permission notice shall be included in all
# copies or substantial portions of the Software.
# 
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.


import os
import re
import sys
import logging
import subprocess
import contextlib
import pathlib

import numpy as np

from .environ import *


@contextlib.contextmanager
def workdir(destination):
    original_dir = pathlib.Path.cwd()
    try:
        os.chdir(destination)
        logging.info(f"Changing working directory to {destination}")
        yield original_dir
    finally:
        os.chdir(original_dir)
        logging.info(f"Go back to {original_dir}")

def abort(err_msg, exit_code=1):
    logging.error(f"{err_msg}. Aborting...")
    sys.exit(exit_code)

def cmd(command, hide=False):
    if not hide:
        logging.info(f"> {command}")

    try:
        result = subprocess.run(command, shell=True, check=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if (disp := result.stdout.decode('utf-8', errors='replace')):
            logging.debug(f"Command output:\n{disp}")
    except subprocess.CalledProcessError as e:
        logging.error(f"Command '{command}' failed with return code: {e.returncode}")
        if e.output:
            logging.error(f"Output: {e.output.decode('utf-8', errors='replace')}")
        if e.stderr:
            logging.error(f"Error output: {e.stderr.decode('utf-8', errors='replace')}")
        raise

def cleandir(directory, abort_on_exist=False, copyfrom=None, ext=None) -> pathlib.Path:
    fdir = pathlib.Path(directory)
    if ext and (fdir.suffix != f".{ext}"):
        fdir = fdir.with_suffix(f".{ext}")

    # clear target directory if exists
    if fdir.exists() and fdir.is_dir():
        if abort_on_exist:
            abort(f"Directory {fdir} exists")
        else:
            logging.debug(f"Directory {fdir} exists. Clearing contents for new files.")
            cmd(f"rm -rf {fdir}")

    # create target directory
    if copyfrom:
        copyfrom = pathlib.Path(copyfrom)
        if ext and (copyfrom.suffix != f".{ext}"):
            copyfrom = copyfrom.with_suffix(f".{ext}")
        logging.debug(f"Copy {copyfrom} to {fdir}.")
        cmd(f"cp -r {copyfrom} {fdir}")
    else:
        logging.debug(f"Created directory {fdir} for new files.")
        cmd(f"mkdir {fdir}")
    return fdir

def checkdir(directory_or_file, abort_msg=None, ext=None) -> pathlib.Path:
    fdir = pathlib.Path(directory_or_file)
    if ext and (fdir.suffix != f".{ext}"):
        fdir = fdir.with_suffix(f".{ext}")

    if not fdir.exists():
        if not abort_msg:
            fdir_type = "Directory" if fdir.is_dir() else "File"
            abort_msg = f"{fdir_type} {fdir} does not exist"
        abort(abort_msg)
    return fdir

def read_dimmat():
    if os.path.exists(PHONOKIT_DIMMAT):
        logging.debug(f"Reading dimensions of supercell from file:  {PHONOKIT_DIMMAT}")
        with open(PHONOKIT_DIMMAT, "r") as f:
            dimstr = f.readline().strip()
    elif os.path.exists(PHONOKIT_DISPYAML):
        logging.debug(f"Reading dimensions of supercell from file:  {PHONOKIT_DISPYAML}")
        with open(PHONOKIT_DISPYAML, "r") as f:
            for line in f:
                if line.startswith("supercell_matrix"):
                    break
            else:
                raise RuntimeError("Cannot find keyword 'supercell_matrix' in {PHONOKIT_DISPYAML}.")
            dimmat = []
            for _ in range(3):
                dimmat.extend(re.findall(r"\-?\d+", f.readline()))
            dimstr = " ".join(dimmat)
    else:
        raise RuntimeError("Cannot find {PHONOKIT_DIMMAT} or {PHONOKIT_DISPYAML}.")
    # logging.debug(f"Dimensions:  {dimstr}")
    return dimstr

def write_dimmat(dimmat, abort_on_exist=False):
    if isinstance(dimmat, str):
        dimstr = dimmat.strip()
    else:
        dimmat = np.asarray(dimmat)
        dimstr = " ".join(f"{s:.0f}" for s in dimmat.T.flatten())

    if os.path.exists(PHONOKIT_DIMMAT):
        if abort_on_exist:
            abort(f"File {PHONOKIT_DIMMAT} exists")
        else:
            logging.warning(f"File {PHONOKIT_DIMMAT} exists. Overwrite...")

    with open(PHONOKIT_DIMMAT, "w") as f:
        f.write(dimstr)
    # logging.debug(f"Write to {PHONOKIT_DIMMAT}:  {dimstr}")

def cubicize_dimmat(lattice, natom, nref=100, nmin=None, nmax=None):
    scale = nref / natom
    det = np.linalg.det(lattice)
    rec = np.linalg.inv(lattice)
    dimmat = np.round(rec * np.power(scale * det, 1/3))
    nreal = natom * np.linalg.det(dimmat)
    dimstr = " ".join(f"{s:.0f}" for s in dimmat.T.flatten())
    logging.info(f"Cubicizing the supercell with dimensions:  {dimstr} (natom: {nreal:.0f})")

    if nmin and nreal < nmin:
        abort(f"The number of atoms in the supercell is too small: {nreal:.0f} < {nmin}")
    if nmax and nreal > nmax:
        abort(f"The number of atoms in the supercell is too large: {nreal:.0f} > {nmax}")
    return dimmat

def read_forcetype():
    if os.path.exists("FORCE_SETS"):
        return "FORCE_SETS"
    elif os.path.exists("FORCE_CONSTANTS"):
        return "FORCE_CONSTANTS"
    raise RuntimeError("Cannot find FORCE_SETS or FORCE_CONSTANTS.")

def read_poscar(poscar, skip_coordinates=False):
    posdict = dict()
    with open(poscar, "r") as f:
        # header
        posdict["header"] = f.readline().strip()
        logging.debug(f"Header:  {posdict['header']}")

        # factor
        if (factor := float(f.readline())) < 0:
            logging.error(f"Negative factor:  {factor}")
            raise NotImplementedError("Negative factor is not supported yet.")
        posdict["factor"] = factor
        logging.debug(f"Factor:  {factor}")

        # lattice
        lattice = []
        for _ in range(3):
            lattice.append([float(s) for s in f.readline().strip().split()])
        posdict["lattice"] = np.array(lattice)
        logging.debug(f"Lattice:  {np.array_str(posdict['lattice'].flatten(), precision=2)}")

        # species
        posdict["species"] = [s.strip() for s in f.readline().strip().split()]
        logging.debug(f"Species:  {posdict['species']}")

        # numbers
        posdict["numbers"] = [int(s) for s in f.readline().strip().split()]
        if len(posdict["species"]) != len(posdict["numbers"]):
            raise ValueError("Unexpected error: Number of species and number of atoms do not match.")
        logging.debug(f"Numbers:  {posdict['numbers']}")

        if skip_coordinates:
            return posdict

        # TODO: check dynamic tag
        # TODO: check selective
        # TODO: check cartesian
        # TODO: check vasp4 version
        raise NotImplementedError

    return posdict

def write_poscar():
    raise NotImplementedError

def get_kpath(path_to_poscar, exec_vaspkit):
    condict = dict()
    cmd(f"mkdir -p {PHONOKIT_TMPFILE} && cp {path_to_poscar} {PHONOKIT_TMPFILE}")
    with workdir(PHONOKIT_TMPFILE):
        cmd(f"(echo 3; echo 305; echo 3) | {exec_vaspkit}")
        with open("KPATH.phonopy", "r") as f:
            for rawline in f:
                if (line := rawline.strip()):
                    key, value = line.strip().split("=", maxsplit=1)
                    condict[key.strip().lower()] = value.strip()
    cmd(f"rm -rf {PHONOKIT_TMPFILE}")
    if not (band := condict.get("band", "")):
        raise RuntimeError("Unexpected missing required key: 'BAND'.")
    if not (band_labels := condict.get("band_labels", "")):
        raise RuntimeError("Unexpected missing required key: 'BAND_LABELS'.")
    logging.debug(f"Generated k-path:  {band_labels}")
    return band, band_labels

def parse_vg(meshfile="mesh.yaml", output="vg.txt", with_q=True):
    dsp = re.compile(r"-?[\d.]+")
    with open(meshfile, "r") as f, open(output, "w") as fout:
        fout.write("#Freq(THz)")    # 10 characters
        if with_q:
            fout.write("    qa(1)    qb(1)    qc(1)")   # 9 characters per item
        fout.write("  vx(km/s)  vy(km/s)  vz(km/s)  vg(km/s)  weight\n") # 10 characters

        for line in f:
            if "q-position" in line:
                q_position = [float(x) for x in dsp.findall(line)]
            elif "weight" in line:
                weight = float(dsp.findall(line)[0])
            elif "frequency" in line:
                freq = float(dsp.findall(line)[0])
            elif "group_velocity" in line:
                vel = [float(x)/10 for x in dsp.findall(line)]  # A*THz to km/s
                vg = sum(vi*vi for vi in vel) ** 0.5

                fout.write(f"{freq:<10.4f}")
                if with_q:
                    for q_pos in q_position:
                        fout.write(f"{q_pos:>9.4f}")
                for vi in vel:
                    fout.write(f"{vi:>10.4f}")
                fout.write(f"{vg:>10.4f}  {weight:<g}\n")

def parse_klabeles(bandfile="band.yaml", output="band_labels.txt", legacy=False):
    msg = f"Unexpected error: Cannot find '{{}}' in {bandfile}."

    with open(bandfile, "r") as f, open(output, "w") as fout:
        # #####  Raw text :
        # nqpoint: 306
        # npath: 6
        # segment_nqpoint:
        # - 51
        # - 51
        # - 51
        # - 51
        # - 51
        # - 51
        # labels:
        # - [ 'G', 'X' ]
        # - [ 'X', 'M' ]
        # - [ 'M', 'G' ]
        # ...

        # parse: npath
        for line in f:
            if "npath" in line:
                npath = int(line.strip().split()[-1])
                break
        else:
            raise ValueError(msg.format("npath"))

        # parse: segment_nqpoint
        for line in f:
            if "segment_nqpoint" in line:
                break
        else:
            raise ValueError(msg.format("segment_nqpoint"))

        segment_acc = [0,]
        for _, line in zip(range(npath), f):
            seg = int(line.strip().split()[-1])
            segment_acc.append(seg + segment_acc[-1])
        segment_acc[0] = 1

        # parse: labels
        for line in f:
            if "labels" in line:
                break
        else:
            raise ValueError(msg.format("labels"))

        seglabels = []
        last_label = ""
        for _, line in zip(range(npath), f):
            _, k_begin, _, k_end, *_ = line.split("'")
            if not legacy:
                if "G" in k_begin:
                    k_begin = r'\g(G)'
                if "G" in k_end:
                    k_end = r'\g(G)'
                k_begin = re.sub(r"_(\d+)", r"\\-(\1)", k_begin)
                k_end = re.sub(r"_(\d+)", r"\\-(\1)", k_end)
            if last_label:
                if k_begin != last_label:
                    seglabels[-1] = "|".join([last_label, k_begin])
            else:
                seglabels.append(k_begin)
            seglabels.append(k_end)
            last_label = k_end

        # parse: k-distance
        ikpt = 0
        for iseg, label in zip(segment_acc, seglabels):
            for line in f:
                if "distance" not in line:
                    continue
                ikpt += 1

                if ikpt == iseg:
                    if legacy:
                        fout.write(f"{label:15s}{line.strip().split()[-1]}\n")
                    else:
                        fout.write(f"{label:15s}{line.strip().split()[-1][:6]}\n")
                    break

def parse_band(bandfile="band.yaml", output="band.txt"):
    dsp = re.compile(r"-?[\d.]+")

    with open(bandfile, "r") as f, open(output, "w") as fout:
        # ####  Raw text :
        # phonon:
        # - q-position: [    0.0000000,    0.0000000,    0.0000000 ]
        #   distance:    0.0000000
        #   band:
        #   - # 1
        #     frequency:   -0.0007269658
        #     eigenvector:
        #     - # atom 1
        #       - [ -0.00000000004504,  0.00000000000000 ]
        #       - [  0.00000000076363,  0.00000000000000 ]
        #       - [  0.49390385172795,  0.00000000000000 ]
        #     - # atom 2
        #       - [ -0.00000000001811,  0.00000000000000 ]
        # ...

        for line in f:
            if "natom" in line:
                natom = int(dsp.findall(line)[0])
                break
        else:
            raise ValueError(f"Unexpected error: Cannot find 'natom' in {bandfile}.")

        # # parse: mass
        # masses = []
        # for _ in range(natom):
        #     for line in f:
        #         if "mass" in line:
        #             masses.append(float(dsp.findall(line)[0]))
        #             break
        #     else:
        #         raise ValueError(f"Unexpected error: Cannot find 'mass' in {bandfile}.")

        # header: #kpath branch-1 branch-2 branch-3 ...
        fout.write("#kpath")
        for i in range(3*natom):
            fout.write(f" branch-{i+1}")

        for line in f:
            if "distance" in line:
                fout.write(f"\n{dsp.findall(line)[0]}")
            elif "frequency" in line:
                fout.write(f" {dsp.findall(line)[0]}")
            # elif "eigenvector" in line:
            #     for mass in masses:
            #         f.readline()    # skip atom index line
            #         for _ in range(3):
            #             vect = [float(x) for x in dsp.findall(f.readline())]
